import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append('/root/ptv3/PointTransformerV3/Pointcept/pointcept/utils')
from visualization import *
import open3d as o3d
import numpy as np


# ScanNet Meta data
VALID_CLASS_IDS_20 = (
    1,
    2,
    3,
    4,
    5,
    6,
    7,
    8,
    9,
    10,
    11,
    12,
    14,
    16,
    24,
    28,
    33,
    34,
    36,
    39,
)


CLASS_LABELS_20 = (
    "wall",
    "floor",
    "cabinet",
    "bed",
    "chair",
    "sofa",
    "table",
    "door",
    "window",
    "bookshelf",
    "picture",
    "counter",
    "desk",
    "curtain",
    "refrigerator",
    "shower curtain",
    "toilet",
    "sink",
    "bathtub",
    "otherfurniture",
)

SCANNET_COLOR_MAP_20 = {
    0: (0.0, 0.0, 0.0),
    1: (173, 199, 232),
    2: (150, 223, 136),
    3: (31, 120, 180),
    4: (255, 188, 120),
    5: (188, 189, 33),
    6: (139, 85, 73),
    7: (255, 151, 150),
    8: (212, 39, 40),
    9: (196, 176, 212),
    10: (147, 103, 188),
    11: (196, 156, 146),
    12: (22, 189, 209),
    14: (248, 183, 210),
    16: (219, 219, 141),
    24: (255, 126, 14),
    28: (159, 219, 228),
    33: (42, 159, 44),
    34: (111, 128, 145),
    36: (226, 119, 194),
    39: (82, 82, 162),
}


def get_npy_files(folder_path):
    """
    获取指定文件夹下所有.npy文件的路径列表
    
    参数:
        folder_path (str): 目标文件夹路径
    
    返回:
        list: 包含所有.npy文件路径的列表
    """
    npy_files = []
    
    # 检查文件夹是否存在
    if not os.path.exists(folder_path):
        print(f"错误: 文件夹 {folder_path} 不存在")
        return npy_files
    
    # 遍历文件夹及其子文件夹
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            # 检查文件扩展名是否为.npy
            if file.endswith('.npy'):
                npy_files.append(os.path.join(root, file))
    
    return npy_files

def save_point_cloud(coord, color, filename):
    """
    Save point cloud data to a PLY file.
    
    Parameters:
    - coord: numpy array of shape (N, 3) representing the coordinates of the points.
    - color: numpy array of shape (N, 3) representing the colors of the points.
    - filename: string, the name of the file to save the point cloud.
    """
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(coord)
    pcd.colors = o3d.utility.Vector3dVector(color)
    o3d.io.write_point_cloud(filename, pcd)

def base_test():
    coord = np.load("/root/ptv3/PointTransformerV3/Pointcept/data/scannet/val/scene0011_00/coord.npy")
    # print(coord.shape)
    color = np.load("/root/ptv3/PointTransformerV3/Pointcept/data/scannet/val/scene0011_00/color.npy")
    color = color.astype(np.float64)/255  # 确保颜色是 float64 类型
    segment = np.load("/root/ptv3/PointTransformerV3/Pointcept/data/scannet/val/scene0011_00/segment20.npy")
    pred = np.load("/root/ptv3/PointTransformerV3/Pointcept/exp/scannet/semseg-pt-v3m1-0-base/result/scene0011_00_pred.npy")

    color_seg = np.array([
        SCANNET_COLOR_MAP_20[VALID_CLASS_IDS_20[int(i)]] 
        for i in segment.flat
    ], dtype=np.float64) / 255.0  # 确保颜色是 float64 类型
    color_pred = np.array([
        SCANNET_COLOR_MAP_20[VALID_CLASS_IDS_20[int(i)]] 
        for i in pred.flat
    ], dtype=np.float64) / 255.0
    # print(color.shape, color.dtype)  # 应输出类似 (N,3) float64
    # print(set(segment.flatten()))  # 打印前10个点的颜色

    save_point_cloud(coord, color, "color.ply")
    save_point_cloud(coord, color_pred, "pred.ply")
    save_point_cloud(coord, color_seg, "seg.ply")


def visual_process():
    pred_folder = "/root/ptv3/PointTransformerV3/Pointcept/exp/scannet/semseg-pt-v3m1-0-base/result"
    pred_files = get_npy_files(pred_folder)
    for pred_file in pred_files:
        scene = pred_file[81:81+12]

        coord_file = f"/root/ptv3/PointTransformerV3/Pointcept/data/scannet/val/{scene}/coord.npy"
        # segment_file = f"/root/ptv3/PointTransformerV3/Pointcept/data/s3dis/{area}/{room}/segment.npy"

        if not os.path.exists(coord_file):
            print(f"Missing files for {scene}, skipping...")
            continue

        coord = np.load(coord_file)
        # segment = np.load(segment_file)
        pred = np.load(pred_file)

        # color_seg = np.array([S3DIS_COLOR_MAP_20[i.item()] for i in segment], dtype=np.float64) / 255.0
        color_pred = np.array([SCANNET_COLOR_MAP_20[VALID_CLASS_IDS_20[int(i)]] for i in pred.flat], dtype=np.float64) / 255.0

        save_point_cloud(coord, color_pred, f"scannet_val_ply/{scene}_pred.ply")
        # save_point_cloud(coord, color_seg, f"{area}_{room}_seg.ply")


if __name__ == "__main__":
    # save_point_cloud(coord, color, "test.ply")
    # base_test()
    visual_process()
    print("Point cloud saved successfully.")